import cv2
import numpy as np
from skimage.morphology import medial_axis, skeletonize
import math
from PyQt5.QtWidgets import QMessageBox
from save import enable_save_button


def measure_insect_length(main_window):
    """测量昆虫体长并显示结果"""
    # 检查是否有分割好的图像
    if not hasattr(main_window, 'mask') or main_window.mask is None:
        QMessageBox.warning(main_window, "警告", "请先使用SAM工具分割图像")
        return

    # 检查是否设置了比例尺
    try:
        pixel_length = float(main_window.txtUnitLength.text())
        scale_value = float(main_window.txtUnitPixels.text())
        if pixel_length <= 0 or scale_value <= 0:
            raise ValueError("比例尺值必须为正数")
    except ValueError as e:
        QMessageBox.warning(main_window, "警告", f"请正确设置比例尺: {str(e)}")
        return

    # 获取单位
    unit = "μm"
    if main_window.radioMM.isChecked():
        unit = "mm"
    elif main_window.radioM.isChecked():
        unit = "m"

    # 获取比例尺转换系数：每像素代表的实际长度
    scale_ratio = scale_value / pixel_length  # 实际单位/像素

    # 使用掩码提取昆虫
    mask = main_window.mask

    # 计算面积（像素和实际面积）
    pixel_area = calculate_area(mask)
    actual_area = pixel_area * (scale_ratio ** 2)  # 平方是因为面积是二维的

    # 计算骨架
    skeleton = extract_skeleton(mask)

    # 计算体长曲线并获取长度（像素）
    length_pixels, optimized_curve = calculate_curve_length(skeleton)

    # 转换为实际单位的长度
    actual_length = length_pixels * scale_ratio

    # 显示结果
    display_measurement_result(main_window, length_pixels, actual_length, pixel_area, actual_area, unit)

    # 在单个目标对象上绘制骨架和体长曲线
    display_skeleton_on_object(main_window, skeleton, optimized_curve)


def calculate_area(mask):
    """计算掩码的面积（像素数）"""
    return np.sum(mask)


def extract_skeleton(mask):
    """提取二值图像的骨架"""
    # 确保输入是二值图像
    binary_mask = mask.astype(np.bool_)

    # 使用skeletonize算法提取更纤细的骨架
    skeleton = skeletonize(binary_mask)

    return skeleton


def calculate_curve_length(skeleton):
    """计算体长曲线的长度，并使用曼哈顿和欧氏距离优化曲线"""
    # 将骨架转换为坐标点
    points = np.where(skeleton)
    points = np.column_stack((points[1], points[0]))  # x, y 坐标

    if len(points) < 2:
        return 0, []  # 没有足够的点来计算长度

    # 找到骨架的端点
    endpoints = find_endpoints(skeleton)

    if len(endpoints) < 2:
        return 0, []  # 没有足够的端点

    # 选择最远的两个端点
    start_point, end_point = find_farthest_endpoints(endpoints)

    # 使用曼哈顿距离和欧氏距离优化路径
    optimized_curve = optimize_curve(skeleton, start_point, end_point)

    # 计算优化后的曲线长度
    length = calculate_path_length(optimized_curve)

    return length, optimized_curve


def find_endpoints(skeleton):
    """找到骨架的端点"""
    # 创建一个3x3的结构元素
    kernel = np.ones((3, 3), np.uint8)

    # 计算每个像素周围的邻居数量
    neighbors = cv2.filter2D((skeleton > 0).astype(np.uint8), -1, kernel) - 1

    # 在骨架中，端点是只有一个邻居的点
    endpoints = np.logical_and(skeleton, neighbors == 1)

    # 将端点转换为坐标列表
    endpoint_coords = np.where(endpoints)
    endpoint_coords = np.column_stack((endpoint_coords[1], endpoint_coords[0]))  # x, y 坐标

    return endpoint_coords


def find_farthest_endpoints(endpoints):
    """找到距离最远的两个端点"""
    max_distance = 0
    farthest_pair = None

    for i in range(len(endpoints)):
        for j in range(i + 1, len(endpoints)):
            # 计算欧氏距离
            dist = np.sqrt(np.sum((endpoints[i] - endpoints[j]) ** 2))
            if dist > max_distance:
                max_distance = dist
                farthest_pair = (endpoints[i], endpoints[j])

    if farthest_pair is None and len(endpoints) >= 2:
        # 如果没有找到点对，返回前两个点
        return endpoints[0], endpoints[-1]
    elif farthest_pair is None:
        # 如果只有一个端点，创建一个虚拟的第二个点
        p = endpoints[0]
        return p, np.array([p[0] + 1, p[1] + 1])

    return farthest_pair


def manhattan_distance(p1, p2):
    """计算曼哈顿距离"""
    return abs(p1[0] - p2[0]) + abs(p1[1] - p2[1])


def euclidean_distance(p1, p2):
    """计算欧氏距离"""
    return np.sqrt((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2)


def optimize_curve(skeleton, start_point, end_point):
    """使用曼哈顿距离和欧氏距离优化体长曲线"""
    # 构建图形表示骨架
    from scipy.sparse import csr_matrix
    from scipy.sparse.csgraph import shortest_path

    # 将骨架转换为坐标集
    skel_points = np.where(skeleton)
    coords = np.column_stack((skel_points[1], skel_points[0]))  # x, y 坐标

    # 创建坐标到索引的映射
    coord_to_idx = {(x, y): i for i, (x, y) in enumerate(map(tuple, coords))}

    # 找到起点和终点在坐标集中的索引
    start_idx = -1
    end_idx = -1
    min_start_dist = float('inf')
    min_end_dist = float('inf')

    for i, (x, y) in enumerate(coords):
        # 找到最接近起点的骨架点
        start_dist = euclidean_distance((x, y), start_point)
        if start_dist < min_start_dist:
            min_start_dist = start_dist
            start_idx = i

        # 找到最接近终点的骨架点
        end_dist = euclidean_distance((x, y), end_point)
        if end_dist < min_end_dist:
            min_end_dist = end_dist
            end_idx = i

    if start_idx == -1 or end_idx == -1:
        return np.array([start_point, end_point])

    # 构建邻接矩阵
    n = len(coords)
    adjacency = np.zeros((n, n), dtype=np.float32)

    # 为每个骨架点找到相邻点
    for i, (x, y) in enumerate(coords):
        for dx in [-1, 0, 1]:
            for dy in [-1, 0, 1]:
                if dx == 0 and dy == 0:
                    continue

                nx, ny = x + dx, y + dy
                neighbor_idx = coord_to_idx.get((nx, ny), -1)

                if neighbor_idx != -1:
                    # 使用欧氏距离作为边权重
                    weight = np.sqrt(dx * dx + dy * dy)
                    adjacency[i, neighbor_idx] = weight

    # 使用Dijkstra算法找到最短路径
    graph = csr_matrix(adjacency)
    dist_matrix, predecessors = shortest_path(graph, directed=False,
                                              indices=start_idx,
                                              return_predecessors=True)

    # 重建从起点到终点的路径
    path = []
    current = end_idx

    while current != start_idx:
        path.append(coords[current])
        current = predecessors[current]

    path.append(coords[start_idx])
    path.reverse()

    # 添加实际的起点和终点
    final_path = [start_point]
    final_path.extend(path)
    final_path.append(end_point)

    return np.array(final_path)


def calculate_path_length(path):
    """计算路径长度"""
    length = 0
    for i in range(1, len(path)):
        # 计算欧氏距离
        length += euclidean_distance(path[i - 1], path[i])

    return length


def display_measurement_result(main_window, length_pixels, actual_length, pixel_area, actual_area, unit):
    """显示测量结果"""
    # 格式化输出结果
    if actual_length < 1 and unit != "μm":
        # 自动转换到更小的单位
        if unit == "m":
            actual_length *= 1000
            actual_area *= 1000000  # 平方米转平方毫米
            unit = "mm"
        elif unit == "mm":
            actual_length *= 1000
            actual_area *= 1000000  # 平方毫米转平方微米
            unit = "μm"

    # 确定面积单位
    area_unit = unit + "²"

    # 格式化显示，保留2位小数
    result_text = f"体长（像素）: {length_pixels:.2f} 像素\n"
    result_text += f"体长（实际）: {actual_length:.2f} {unit}\n\n"
    result_text += f"面积（像素）: {pixel_area} 像素²\n"
    result_text += f"面积（实际）: {actual_area:.2f} {area_unit}\n"

    # 将结果显示在文本框中
    main_window.txtResults.setText(result_text)

    # 更新状态栏
    main_window.statusbar.showMessage(
        f"测量完成: 昆虫体长为 {actual_length:.2f} {unit}, 面积为 {actual_area:.2f} {area_unit}")

    # 启用保存按钮
    enable_save_button(main_window)


def display_skeleton_on_object(main_window, skeleton, optimized_curve):
    """在单个目标对象上显示骨架和优化后的体长曲线"""
    if main_window.source_image is None or main_window.mask is None:
        return

    # 创建一个只包含目标对象的图像
    object_img = np.zeros_like(main_window.source_image)
    object_img[main_window.mask] = main_window.source_image[main_window.mask]

    # 创建一个用于显示的彩色图像
    display_img = object_img.copy()

    # 在目标对象上绘制骨架（黄色）
    skeleton_coords = np.where(skeleton)
    for y, x in zip(skeleton_coords[0], skeleton_coords[1]):
        if main_window.mask[y, x]:  # 只绘制在目标对象内部的骨架点
            cv2.circle(display_img, (x, y), 1, (0, 255, 255), -1)

    # 在目标对象上绘制优化后的体长曲线（蓝色）
    if optimized_curve is not None and len(optimized_curve) > 1:
        for i in range(len(optimized_curve) - 1):
            pt1 = tuple(optimized_curve[i].astype(int))
            pt2 = tuple(optimized_curve[i + 1].astype(int))
            cv2.line(display_img, pt1, pt2, (255, 0, 0), 2)

        # 在曲线的端点上绘制大一点的标记
        cv2.circle(display_img, tuple(optimized_curve[0].astype(int)), 5, (0, 0, 255), -1)
        cv2.circle(display_img, tuple(optimized_curve[-1].astype(int)), 5, (0, 0, 255), -1)

    # 显示处理后的图像在单个对象框中
    main_window.display_single_object(main_window.mask)  # 确保显示原始目标对象

    # 更新分割图像显示结果
    main_window.segmented_image = display_img
    main_window.display_segmented_image(display_img)


def start_measurement(main_window):
    """启动测量过程"""
    # 检查是否有图像和掩码
    if main_window.source_image is None:
        QMessageBox.warning(main_window, "警告", "请先打开图像")
        return

    if not hasattr(main_window, 'mask') or main_window.mask is None:
        QMessageBox.warning(main_window, "警告", "请先使用SAM工具分割图像")
        return

    # 执行测量
    measure_insect_length(main_window)
